package com.example.origin.finance.service;

import cn.hutool.core.util.HashUtil;
import cn.hutool.core.util.ObjectUtil;
import com.example.origin.finance.dto.BaseTable;
import com.example.origin.finance.enums.TableNameEnum;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.util.CollectionUtils;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public abstract class AbstractMultiStoreService<T extends BaseTable> extends AbstractUtilService {


   private static class DefaultLinkedMultiValueMap<String, T> extends LinkedMultiValueMap<String, T> {
       private int size;
       private Set<String> codes;
       public DefaultLinkedMultiValueMap(int expectedSize) {
           super(expectedSize);
           this.size = expectedSize;
           this.codes=new HashSet<>(expectedSize);
       }

        @Override
        public synchronized void add(String key, T value) {
            super.add(key, value);
            if (!codes.contains(key)) {
                codes.add(key);
                logProcess(this.codes.size(), size);
            }
        }
    }

    @SneakyThrows
    public void storeTable(List<String> list, TableNameEnum tableNameEnum) {
        String name= tableNameEnum.getTable();
        Integer tableCount = tableNameEnum.getCount();
        boolean onlyOne = tableCount == 1;
        String table = createTableName(name, 0, onlyOne);
        int size = list.size();
        boolean flag = size == 1;
        if (!flag) {
            for (int i = 0; i < tableCount; i++) {
                if (mongoTemplate.collectionExists(createTableName(name,i,onlyOne))) {
                    mongoTemplate.dropCollection(createTableName(name,i,onlyOne));
                }
            }
        }

        if (flag || !mongoTemplate.collectionExists(table)) {

            ExecutorService dataGetExecutorService = Executors.newFixedThreadPool(20);
            MultiValueMap<String, T> linkedMultiValueMap = new DefaultLinkedMultiValueMap<>(size);
            List<CompletableFuture<Void>> futureList = new ArrayList<>();
            for (int i = 0; i < size; i++) {
                String code = list.get(i);
                CompletableFuture<Void> runAsync = CompletableFuture.runAsync(() -> {
                        fillMap(code, linkedMultiValueMap);
                }, dataGetExecutorService);
                futureList.add(runAsync);
            }
            CompletableFuture.allOf(futureList.toArray(new CompletableFuture[0])).join();
            dataGetExecutorService.shutdown();
            boolean a;
            do {
                a = dataGetExecutorService.awaitTermination(1, TimeUnit.SECONDS);
            } while (!a);

            MultiValueMap<String, T> klineMultiValueMap = multiMap(linkedMultiValueMap, name, tableCount);
            if (klineMultiValueMap.isEmpty()) {
                return;
            }
            CountDownLatch countDownLatch = new CountDownLatch(klineMultiValueMap.size());
            ExecutorService executorService = Executors.newFixedThreadPool(100);
            for (Map.Entry<String, List<T>> listEntry : klineMultiValueMap.entrySet()) {
                executorService.execute(() -> {
                    String tableName = listEntry.getKey();
                    List<T> value = listEntry.getValue();
                    try {
                        if (CollectionUtils.isEmpty(value)) {
                            return;
                        }
                        log.info("start >> tableName：{}，count：{}", tableName, value.size());
                        if (flag && mongoTemplate.collectionExists(tableName)) {
                            mongoTemplate.remove(Query.query(Criteria.where("code").is(value.get(0).getCode())), tableName);
                        }
                        mongoTemplate.insert(value, tableName);
                        if (!flag) {
                            createIndex(tableName);
                        }
                    } finally {
                        countDownLatch.countDown();
                        log.info("end >> tableName：{}，count：{}", tableName, value.size());
                    }
                });
            }
            countDownLatch.await();
            executorService.shutdown();
            boolean b;
            do {
                b = executorService.awaitTermination(1, TimeUnit.SECONDS);
            } while (!b);
        }
    }

    private MultiValueMap<String, T> multiMap(MultiValueMap<String, T> linkedMultiValueMap, String name, Integer tableNum) {
        Set<String> codes = linkedMultiValueMap.keySet();

        MultiValueMap<Integer, String> multiValueMap = new LinkedMultiValueMap<>();
        for (String code : codes) {
            multiValueMap.add(HashUtil.rotatingHash(code, tableNum), code);
        }
        MultiValueMap<String, T> valueMap = new LinkedMultiValueMap<>();
        for (int i = 0; i < tableNum; i++) {
            if (multiValueMap.containsKey(i)) {
                for (String s : Objects.requireNonNull(multiValueMap.get(i))) {
                    List<T> klineList = linkedMultiValueMap.get(s);
                    if (!ObjectUtil.isEmpty(klineList)) {
                        valueMap.addAll(createTableName(name, i,tableNum==1), Objects.requireNonNull(klineList));
                    }
                }
            }
        }
        return valueMap;
    }

    protected abstract void fillMap(String code, MultiValueMap<String, T> linkedMultiValueMap);

    protected static String createTableName(String s, int i,boolean onlyOne) {
        if (onlyOne){
            return s;
        }
        return s + "_" + i;
    }

}
